import bisect
import copy
import math
import string

import matplotlib
from matplotlib import pyplot
import numpy

from fitter import Fitter
from histogram import Histogram1D

class BinnedFitter(Fitter):
    '''
    Provides a nicer interface for fitting histograms with the scipy.optimize algorithms. Inherits
    these attributes from ``Fitter``:

    Inherited attributes
    --------------------------

    x             -- X-coordinates of the data to be fit
    y             -- Y-coordinates of the data to be fit
    stats         -- ``FitStatistics`` object containing the results of the last fit
    par_names     -- A list containing the names of the fit parameters
    npar          -- The number of parameters used in the fit
    fixedp        -- A list of indices of parameters that are not allowed to vary during a fit
    fixed_bg_pars -- A list of parameters values to supply to a fixed background function, if any.
    data_len      -- The number of points to be fit

    Native Attributes
    --------------------------

    hist          -- the histogram object responsible for binning the data
    data          -- the unbinned data itself, if provided

    '''
    
    fcn_str = '1'
    def __init__(self, finput, par_names, fcn_str='1'):
        '''
        Default constructor. The end user should not call this directly; instead he or she should use one
        of the derived BinnedFitters that implement a specific fit function.

        Arguments:

        finput    -- either an iterable containing data to be binned later, or a Histogram1D object
        par_names -- a list of names for the fit parameters
        fcn_str   -- a Python-compilable string representing the fit function

        '''

        self.hist = None
        self.data = None        

        try: 
            it = iter(finput)
        except TypeError:
            self.hist = copy.deepcopy(finput)
            self.data_x_max = self.hist.data_max()
            self.data_x_min = self.hist.data_min()
            self.data_mean  = self.hist.data_mean()
            self.data_std   = self.hist.data_std()
            self.data_len   = self.hist.nentries()
            
        else:
            self.data = finput
            self.data_x_max = max(self.data)
            self.data_x_min = min(self.data)
            self.data_mean  = numpy.mean(self.data)
            self.data_std   = numpy.std(self.data)
            self.data_len   = len(self.data)

        super(BinnedFitter, self).__init__(range(5), range(5), par_names, fcn_str=fcn_str)
        
        if self.hist is not None:
            self.title      = self.hist.title
            self.xlabel     = self.hist.xlabel
            self.ylabel     = self.hist.ylabel
        
    def _hist_data(self, nbins, hrange, loud):
        '''
        Bin the data according to the parameters provided.

        Arguments:

        nbins  -- number of bins of the histogram
        hrange -- a 2-tuple indicating the X-axis range for the histogram
        loud   -- when True, causes the histogram to write messages to stdout as it bins the data
        '''
        
        if hrange is None:
            hrange = self._set_hrange_none()
        elif type(hrange) == type(float(1)) or type(hrange) == type(int(1)):
            hrange = self._set_hrange_float(hrange)
            
        self.hist = Histogram1D('MyFit', nbins, hrange, 'my fitted data', '', bufsize=len(self.data))
        self.hist.fill(self.data, loud)

    def _set_hrange_none(self):
        '''
        Set a default histogram range (+/- 5 sigma about the mean) when one is not provided.
        '''
        return (self.data_mean - 5*self.data_std, self.data_mean + 5*self.data_std)
        
    def _set_hrange_float(self, hrange):
        '''
        Set the histogram range as the mean +/- a multiple of the standard deviation of the data.
        '''
        
        return (self.data_mean - hrange*self.data_std, self.data_mean + hrange*self.data_std)

    def _set_frange_none(self):
        '''
        Set a default fit range when none is specified (+/- 3 sigma about the mean).
        '''
        
        return (self.data_mean - 3*self.data_std, self.data_mean + 3*self.data_std)
        
    def _set_frange_float(self, frange):
        '''
        Set the fit range as the maximum +/- a multiple of the standard deviation of the data.
        '''
        
        xmax = self.hist.max_bin_x()
        return (xmax - frange, xmax + frange)

    def fit(self, nbins=100, initpars=None, hrange=None, frange=None, par_names=None, bins_to_omit=[], loud=False):
        '''
        Perform a fit.

        Keyword arguments:

        nbins        -- Bin the data into ``nbins`` bins
        initpars     -- a list of initial parameters
        hrange       -- a 2-tuple describing the range of the histogram
        frange       -- a 2-tuple describing the fit range
        par_names    -- a list of names for the parameters for this fit
        bins_to_omit -- a list of bin indicies to remove from the fit
        loud         -- if True, if the data must be binned, have the histogram report progress to stdout

        '''

        if self.hist is None:
            self._hist_data(nbins, hrange, loud)
        
        self.hist.flush()

        if frange is None:
            frange = self._set_frange_none()
        elif type(frange) == type(float(1.0)) or type(frange) == type(int(1)):
            frange = self._set_frange_float(frange)

        max_bin_index = self.hist.max_bin_index()

        fit_bin_centers = list()
        fit_bins        = list()
        bins            = self.hist.bins()

        for i in xrange(len(bins)):
            if i in bins_to_omit:
                if i == max_bin_index:
                    bins[i] = 0
                continue

            fit_bin_centers.append(self.hist.low_edges[i] + self.hist.bin_width/2.0)
            fit_bins.append(bins[i])

        self.x = numpy.array(fit_bin_centers)
        self.y = numpy.array(fit_bins)

        return super(BinnedFitter, self).fit(initpars=initpars, frange=frange, par_names=par_names)
        


    def draw(self, name='MyFit', ext='png', histtype='stepfilled', facecolor='#0080FF', plotarg='r-', plotlw=2, 
             hist_kwargs=dict(), plot_kwargs=dict(), pars=None, same=False, hist=True, fit=True, stats=True, hold=False,
             **kwargs):
        '''
        Draw the result of a binned fit.

        Keyword arguments:
        
        name        -- a name for the fit. The output file will have the form ``name.ext``.
        ext         -- the desired image extension. Allowed values are ``png``, ``gif``, and ``jpg``.
        plotarg     -- a Matplotlib-like string defining how the fit function should be drawn.
        plotlw      -- the width of the line of the drawn function
        pars        -- the parameters with which the function should be drawn (defaults to the last fit parameters)
        hist_kwargs -- a dictionary of keyword arguments to be passed to the underlying ``pyplot.hist()`` call
        plot_kwargs -- a dictionary of keyword arguments to be passed to the underlying ``pyplot.plot()`` call
        same        -- draw without clearing the canvas
        hist        -- draw the histogram
        fit         -- draw the resulting fit
        stats       -- write the fit parameters and statistics to the image
        hold        -- do not automatically write to file
        
        '''

        orig_attrs  = dict()
        local_attrs = list()
        for key in kwargs:
            try:
                orig_attrs[key] = getattr(self, key)
            except AttributeError:
                local_attrs.append(key)

            setattr(self, key, kwargs[key])

        if self.hist is None:
            self._hist_data(100, None, True)
            return self.hist.draw(name, ext, histtype, facecolor, **hist_kwargs)
            
        if pars is None:
            pars = self.stats.fitpars
            
        if not same:
            pyplot.clf()
            if stats:
                pyplot.figure(figsize=(14, 8.5))
                pyplot.subplots_adjust(right=0.6)

        if not self.hist.name:
            self.hist.name = name

        if hist:
            self.hist.draw(hold=True, **hist_kwargs)
        
        if fit:
            pyplot.plot(self.stats.x,
                        [self._fcn(pars, x) for x in self.stats.x],
                        plotarg,
                        linewidth=plotlw,
                        **plot_kwargs)

        pyplot.title(self.title)
        pyplot.xlabel(self.xlabel)
        pyplot.ylabel(self.ylabel)

        if stats:
            if matplotlib.rcParams['text.usetex']:
                self._draw_stats_latex(pars)
            else:
                self._draw_stats(pars)

        if not hold:
            if name == 'MyFit' and self.hist.name is not '':
                return self.save_and_clear()
            else:
                return self.save_and_clear(name=name)

    def save_and_clear(self, **kwargs):
        retname = self.hist.save_and_clear(**kwargs)
        pyplot.subplots_adjust(right=0.9)
        return retname

    def _draw_stats_latex(self, pars):
        '''
        Draw the fit parameters and statistics on the current canvas with full LaTeX markup.

        Arguments:

        pars -- The list of parameters to be written out to the canvas.
        '''
        
        npars = len(pars)


#        statstr  = r"\begin{equation*}" + self.fcn_latex[1:-1] + r"\end{equation*}\\" 
        statstr  = self.fcn_latex[:] + r'\\ \vskip 5mm'
        statstr  += r"\begin{tabular}{lll}"

        if self.stats.par_names is not None:
            for i, par in enumerate(pars):
                statstr += r"{0}&= {1:> .3g}&$\pm${2:.3g}\\".format(self.stats.par_names[i], par, 
                                                                      self.stats.errpars[i])
        else:
            for i, par in enumerate(pars):
                statstr += r"$p_{0}$ &= \text{{ {1:> .3g} }} &\pm \text{{ {2:.3g} }}\\".format(i, par, self.stats.errpars[i])
        
        statstr += r" \end{tabular}\vskip 5mm"
        statstr  += r"\begin{tabular}{lr}"
    
        statstr += r"Entries: &{0:1d}\\".format(self.hist.nentries())
        statstr += r"$(\chi^2/d.o.f.)_i$: &{0:.1e}\\".format(self.stats.init_chi2_dof)
        statstr += r"$(\chi^2/d.o.f.)_f$: &{0:.1e}\\".format(self.stats.final_chi2_dof)
        statstr += r"n.d.o.f.: &" + str(self.stats.ndof)


        statstr += r" \end{tabular}"

#        statstr += r"Parameters:"

        pyplot.figtext(0.65, 0.8, statstr, size='x-large', va='top')
        

    def _draw_stats(self, pars):
        '''
        Draw the fit parameters and statistics on the current canvas.

        Arguments:

        pars -- The list of parameters to be written out to the canvas.
        '''
        
        fitstr = ""

        if self.stats.par_names is not None:
            for i in xrange(len(pars)):
                fitstr += "  {0} = {1:.3e}\n".format(self.stats.par_names[i], pars[i])
        else:
            for i in xrange(len(pars)):
                fitstr += "  $p_{0}$ = {1:.3e}\n".format(i, pars[i])

        pars_str = "Parameters:\n"
        pars_str += "  " + self.fcn_latex + '\n\n'
        pars_str += fitstr

        stats_str = "Entries: {0:1d}\n".format(self.hist.nentries())
        stats_str += "$(\\chi^2/d.o.f.)_i$: {0:.1e}\n".format(self.stats.init_chi2_dof)
        stats_str += "$(\\chi^2/d.o.f.)_f$: {0:.1e}\n".format(self.stats.final_chi2_dof)
        stats_str += "n.d.o.f.: " + str(self.stats.ndof) + '\n'

        pyplot.figtext(0.65, self._par_text_offset, pars_str, size='x-large')
        pyplot.figtext(0.65, 0.045, stats_str, size='x-large')

